#!/usr/bin/env python3
"""
高级图表生成函数

使用统一配色方案和样式，生成专业科研论文级别的图表
"""

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from typing import Dict, List, Tuple, Optional, Any, Union
from pathlib import Path

try:
    import torch
except ImportError:
    torch = None

import sys
sys.path.append(str(Path(__file__).parent.parent.parent))
from .color_palette import ColorPalette
from .plot_style import setup_plot_style
from utils.logger import get_logger


class ChartGenerator:
    """图表生成器"""
    
    def __init__(self):
        """初始化图表生成器"""
        self.color_palette = ColorPalette()
        self.style = setup_plot_style()
        self.logger = get_logger("ChartGenerator")
    
    def plot_comparison_bar(self,
                           data: Dict[str, Dict[str, float]],
                           metrics: List[str],
                           output_path: Optional[str] = None,
                           title: str = "Comparison Results",
                           figsize: Tuple[int, int] = (10, 6)) -> plt.Figure:
        """
        绘制对比实验柱状图（带渐变配色）
        
        Args:
            data: 数据字典 {method_name: {metric: value}}
            metrics: 指标列表
            output_path: 输出路径
            title: 图表标题
            
        Returns:
            matplotlib Figure对象
        """
        fig, ax = plt.subplots(figsize=figsize)
        
        methods = list(data.keys())
        n_methods = len(methods)
        n_metrics = len(metrics)
        
        # 获取颜色
        colors = self.color_palette.get_comparison_colors(n_methods)
        
        # 设置x轴位置
        x = np.arange(n_metrics)
        width = 0.8 / n_methods
        
        # 绘制柱状图
        for i, method in enumerate(methods):
            values = [data[method].get(metric, 0) for metric in metrics]
            offset = (i - n_methods / 2 + 0.5) * width
            bars = ax.bar(x + offset, values, width, label=method, 
                         color=colors[i], alpha=0.8, edgecolor='white', linewidth=1)
            
            # 添加数值标签
            for bar in bars:
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width()/2., height,
                       f'{height:.3f}',
                       ha='center', va='bottom', fontsize=9)
        
        # 设置标签和标题
        ax.set_xlabel('Metrics', fontsize=self.style.FONT_SIZE_AXIS)
        ax.set_ylabel('Score', fontsize=self.style.FONT_SIZE_AXIS)
        ax.set_title(title, fontsize=self.style.FONT_SIZE_TITLE, pad=15)
        ax.set_xticks(x)
        ax.set_xticklabels(metrics, rotation=0)
        ax.legend(loc='best', frameon=True, fontsize=self.style.FONT_SIZE_LEGEND)
        ax.grid(True, alpha=0.3, linestyle='--')
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            self.logger.info(f"图表已保存: {output_path}")
        
        return fig
    
    def plot_ablation_heatmap(self,
                             data: np.ndarray,
                             row_labels: List[str],
                             col_labels: List[str],
                             output_path: Optional[str] = None,
                             title: str = "Ablation Study",
                             figsize: Tuple[int, int] = (8, 6)) -> plt.Figure:
        """
        绘制消融实验热力图（专业配色方案）
        
        Args:
            data: 数据矩阵 (n_rows, n_cols)
            row_labels: 行标签
            col_labels: 列标签
            output_path: 输出路径
            title: 图表标题
            
        Returns:
            matplotlib Figure对象
        """
        fig, ax = plt.subplots(figsize=figsize)
        
        # 使用viridis配色
        cmap = self.color_palette.get_heatmap_colormap('viridis')
        
        # 绘制热力图
        im = ax.imshow(data, cmap=cmap, aspect='auto', interpolation='nearest')
        
        # 设置标签
        ax.set_xticks(np.arange(len(col_labels)))
        ax.set_yticks(np.arange(len(row_labels)))
        ax.set_xticklabels(col_labels, fontsize=self.style.FONT_SIZE_AXIS)
        ax.set_yticklabels(row_labels, fontsize=self.style.FONT_SIZE_AXIS)
        
        # 添加数值标注
        for i in range(len(row_labels)):
            for j in range(len(col_labels)):
                text = ax.text(j, i, f'{data[i, j]:.3f}',
                             ha="center", va="center",
                             color="white" if data[i, j] < data.mean() else "black",
                             fontsize=self.style.FONT_SIZE_LABEL)
        
        # 添加颜色条
        cbar = plt.colorbar(im, ax=ax)
        cbar.set_label('Score', fontsize=self.style.FONT_SIZE_AXIS)
        
        ax.set_title(title, fontsize=self.style.FONT_SIZE_TITLE, pad=15)
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            self.logger.info(f"热力图已保存: {output_path}")
        
        return fig
    
    def plot_training_curves(self,
                            train_loss: List[float],
                            val_loss: Optional[List[float]] = None,
                            train_metrics: Optional[Dict[str, List[float]]] = None,
                            val_metrics: Optional[Dict[str, List[float]]] = None,
                            output_path: Optional[str] = None,
                            title: str = "Training Curves",
                            figsize: Tuple[int, int] = (12, 8)) -> plt.Figure:
        """
        绘制训练曲线图（平滑渐变曲线）
        
        Args:
            train_loss: 训练损失列表
            val_loss: 验证损失列表（可选）
            train_metrics: 训练指标字典（可选）
            val_metrics: 验证指标字典（可选）
            output_path: 输出路径
            title: 图表标题
            
        Returns:
            matplotlib Figure对象
        """
        n_plots = 1
        if train_metrics:
            n_plots += len(train_metrics)
        
        fig, axes = plt.subplots(1, n_plots, figsize=figsize)
        if n_plots == 1:
            axes = [axes]
        
        epochs = range(1, len(train_loss) + 1)
        
        # 绘制损失
        ax = axes[0]
        colors = self.color_palette.get_gradient_colors('blue', len(train_loss))
        ax.plot(epochs, train_loss, 'o-', label='Train Loss', 
               color=self.color_palette.BLUE_GRADIENT[0], linewidth=2, markersize=4)
        
        if val_loss:
            ax.plot(epochs, val_loss, 's-', label='Val Loss',
                   color=self.color_palette.ORANGE_GRADIENT[0], linewidth=2, markersize=4)
        
        ax.set_xlabel('Epoch', fontsize=self.style.FONT_SIZE_AXIS)
        ax.set_ylabel('Loss', fontsize=self.style.FONT_SIZE_AXIS)
        ax.set_title('Loss', fontsize=self.style.FONT_SIZE_TITLE)
        ax.legend(fontsize=self.style.FONT_SIZE_LEGEND)
        ax.grid(True, alpha=0.3, linestyle='--')
        
        # 绘制指标
        if train_metrics:
            for idx, (metric_name, metric_values) in enumerate(train_metrics.items(), 1):
                if idx < len(axes):
                    ax = axes[idx]
                    ax.plot(epochs, metric_values, 'o-', label=f'Train {metric_name}',
                           color=self.color_palette.BLUE_GRADIENT[0], linewidth=2, markersize=4)
                    
                    if val_metrics and metric_name in val_metrics:
                        ax.plot(epochs, val_metrics[metric_name], 's-', 
                               label=f'Val {metric_name}',
                               color=self.color_palette.ORANGE_GRADIENT[0], 
                               linewidth=2, markersize=4)
                    
                    ax.set_xlabel('Epoch', fontsize=self.style.FONT_SIZE_AXIS)
                    ax.set_ylabel(metric_name, fontsize=self.style.FONT_SIZE_AXIS)
                    ax.set_title(metric_name, fontsize=self.style.FONT_SIZE_TITLE)
                    ax.legend(fontsize=self.style.FONT_SIZE_LEGEND)
                    ax.grid(True, alpha=0.3, linestyle='--')
        
        fig.suptitle(title, fontsize=self.style.FONT_SIZE_TITLE + 2, y=1.02)
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            self.logger.info(f"训练曲线已保存: {output_path}")
        
        return fig
    
    def plot_correlation_scatter(self,
                                 x: np.ndarray,
                                 y: np.ndarray,
                                 labels: Optional[List[str]] = None,
                                 output_path: Optional[str] = None,
                                 title: str = "Correlation Analysis",
                                 xlabel: str = "Predicted",
                                 ylabel: str = "Ground Truth",
                                 figsize: Tuple[int, int] = (8, 8)) -> plt.Figure:
        """
        绘制相关性散点图（带密度渐变）
        
        Args:
            x: x轴数据
            y: y轴数据
            labels: 数据点标签（可选）
            output_path: 输出路径
            title: 图表标题
            
        Returns:
            matplotlib Figure对象
        """
        fig, ax = plt.subplots(figsize=figsize)
        
        # 计算相关系数
        if torch is not None:
            if isinstance(x, torch.Tensor):
                x = x.cpu().numpy()
            if isinstance(y, torch.Tensor):
                y = y.cpu().numpy()
        
        x = x.flatten()
        y = y.flatten()
        
        corr = np.corrcoef(x, y)[0, 1]
        
        # 绘制散点图（带密度渐变）
        scatter = ax.scatter(x, y, alpha=0.6, s=50, c=np.sqrt(x**2 + y**2),
                           cmap=self.color_palette.get_heatmap_colormap('viridis'),
                           edgecolors='white', linewidth=0.5)
        
        # 添加对角线
        min_val = min(x.min(), y.min())
        max_val = max(x.max(), y.max())
        ax.plot([min_val, max_val], [min_val, max_val], 
               'r--', linewidth=2, alpha=0.5, label='y=x')
        
        # 添加相关系数文本
        ax.text(0.05, 0.95, f'Correlation: {corr:.3f}',
               transform=ax.transAxes, fontsize=self.style.FONT_SIZE_AXIS,
               verticalalignment='top',
               bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
        ax.set_xlabel(xlabel, fontsize=self.style.FONT_SIZE_AXIS)
        ax.set_ylabel(ylabel, fontsize=self.style.FONT_SIZE_AXIS)
        ax.set_title(title, fontsize=self.style.FONT_SIZE_TITLE, pad=15)
        ax.legend(fontsize=self.style.FONT_SIZE_LEGEND)
        ax.grid(True, alpha=0.3, linestyle='--')
        
        # 添加颜色条
        cbar = plt.colorbar(scatter, ax=ax)
        cbar.set_label('Distance', fontsize=self.style.FONT_SIZE_AXIS)
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            self.logger.info(f"散点图已保存: {output_path}")
        
        return fig
    
    def plot_temporal_analysis(self,
                              data: Dict[str, np.ndarray],
                              time_points: Optional[List[str]] = None,
                              output_path: Optional[str] = None,
                              title: str = "Temporal Analysis",
                              figsize: Tuple[int, int] = (12, 6)) -> plt.Figure:
        """
        绘制时序性能分析图（渐变折线图）
        
        Args:
            data: 数据字典 {metric_name: values}
            time_points: 时间点标签（可选）
            output_path: 输出路径
            title: 图表标题
            
        Returns:
            matplotlib Figure对象
        """
        fig, ax = plt.subplots(figsize=figsize)
        
        if time_points is None:
            time_points = [f'Frame {i+1}' for i in range(len(list(data.values())[0]))]
        
        colors = self.color_palette.get_comparison_colors(len(data))
        
        for idx, (metric_name, values) in enumerate(data.items()):
            if torch is not None and isinstance(values, torch.Tensor):
                values = values.cpu().numpy()
            values = values.flatten()
            
            ax.plot(time_points, values, 'o-', label=metric_name,
                   color=colors[idx], linewidth=2, markersize=6, alpha=0.8)
        
        ax.set_xlabel('Time Points', fontsize=self.style.FONT_SIZE_AXIS)
        ax.set_ylabel('Score', fontsize=self.style.FONT_SIZE_AXIS)
        ax.set_title(title, fontsize=self.style.FONT_SIZE_TITLE, pad=15)
        ax.legend(fontsize=self.style.FONT_SIZE_LEGEND)
        ax.grid(True, alpha=0.3, linestyle='--')
        plt.xticks(rotation=45, ha='right')
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            self.logger.info(f"时序分析图已保存: {output_path}")
        
        return fig


# 便捷函数
def plot_comparison_bar(*args, **kwargs):
    """绘制对比实验柱状图"""
    generator = ChartGenerator()
    return generator.plot_comparison_bar(*args, **kwargs)


def plot_ablation_heatmap(*args, **kwargs):
    """绘制消融实验热力图"""
    generator = ChartGenerator()
    return generator.plot_ablation_heatmap(*args, **kwargs)


def plot_training_curves(*args, **kwargs):
    """绘制训练曲线图"""
    generator = ChartGenerator()
    return generator.plot_training_curves(*args, **kwargs)


def plot_correlation_scatter(*args, **kwargs):
    """绘制相关性散点图"""
    generator = ChartGenerator()
    return generator.plot_correlation_scatter(*args, **kwargs)


def plot_temporal_analysis(*args, **kwargs):
    """绘制时序性能分析图"""
    generator = ChartGenerator()
    return generator.plot_temporal_analysis(*args, **kwargs)


def plot_distribution_kde(data: np.ndarray,
                          labels: Optional[List[str]] = None,
                          output_path: Optional[str] = None,
                          title: str = "Distribution Analysis",
                          figsize: Tuple[int, int] = (10, 6)) -> plt.Figure:
    """
    绘制分布分析KDE图
    
    Args:
        data: 数据数组或列表
        labels: 标签列表（可选）
        output_path: 输出路径
        title: 图表标题
        
    Returns:
        matplotlib Figure对象
    """
    generator = ChartGenerator()
    fig, ax = plt.subplots(figsize=figsize)
    
    if isinstance(data, list):
        # 多个分布
        colors = generator.color_palette.get_comparison_colors(len(data))
        for idx, (dist_data, label) in enumerate(zip(data, labels or [])):
            if torch is not None and isinstance(dist_data, torch.Tensor):
                dist_data = dist_data.cpu().numpy()
            dist_data = dist_data.flatten()
            sns.kdeplot(data=dist_data, label=label, ax=ax, 
                       color=colors[idx], linewidth=2)
    else:
        # 单个分布
        if torch is not None and isinstance(data, torch.Tensor):
            data = data.cpu().numpy()
        data = data.flatten()
        sns.kdeplot(data=data, ax=ax, 
                   color=generator.color_palette.BLUE_GRADIENT[0], linewidth=2)
    
    ax.set_xlabel('Value', fontsize=generator.style.FONT_SIZE_AXIS)
    ax.set_ylabel('Density', fontsize=generator.style.FONT_SIZE_AXIS)
    ax.set_title(title, fontsize=generator.style.FONT_SIZE_TITLE, pad=15)
    if labels:
        ax.legend(fontsize=generator.style.FONT_SIZE_LEGEND)
    ax.grid(True, alpha=0.3, linestyle='--')
    
    plt.tight_layout()
    
    if output_path:
        plt.savefig(output_path, dpi=300, bbox_inches='tight')
        generator.logger.info(f"KDE图已保存: {output_path}")
    
    return fig
